package rabbitmq.publisher.confirms;


import com.rabbitmq.client.Channel;
import com.rabbitmq.client.ConfirmListener;
import com.rabbitmq.client.Connection;
import com.rabbitmq.client.ConnectionFactory;
import rabbitmq.constants.Constants;

import java.io.IOException;
import java.util.Collections;
import java.util.SortedSet;
import java.util.TreeSet;


public class PublisherConfirms {
    private static final int MESSAGE_COUNT = 2000;

    public static Connection createConnection() throws Exception {
        ConnectionFactory factory = new ConnectionFactory();
        factory.setHost(Constants.HOST);
        factory.setPort(Constants.PORT);
        factory.setUsername(Constants.USERNAME);
        factory.setPassword(Constants.PASSWORD);
        factory.setVirtualHost(Constants.VIRTUAL_HOST);
        return factory.newConnection();
    }

    public static void main(String[] args) throws Exception {
        // 单独确认
        publishingMessagesIndividually();
        // 批量确认
        publishingMessagesInBatches();
        // 异步确认
        handlingPublisherConfirmsAsynchronously();
    }

    /**
     * 异步确认
     */
    private static void handlingPublisherConfirmsAsynchronously() throws Exception {
        try (Connection connection = createConnection()) {
            // 1.开启信道
            Channel channel = connection.createChannel();
            // 2.设置为confirm模式
            channel.confirmSelect();
            // 3.使用内置交换机，声明队列
            channel.queueDeclare(Constants.PUBLISHER_CONFIRMS_QUEUE3,true,false,false,null);
            // 4.监听confirm
            long startTime = System.currentTimeMillis();
            // 创建有序集合，元素按照自然排序，用来存储未被确认的confirm消息序号
            SortedSet<Long> confirmSetNo = Collections.synchronizedSortedSet(new TreeSet<>());
            channel.addConfirmListener(new ConfirmListener() {
                @Override
                public void handleAck(long deliveryTag, boolean multiple) throws IOException {
                    if (multiple) {
                        // 如果开启了批量确认，则批量确认之后要清空有序集合已经确认好的元素
                        confirmSetNo.headSet(deliveryTag+1).clear();
                    } else {
                        // 如果没有开，那么只需要清除当前的这一个
                        confirmSetNo.remove(deliveryTag);
                    }
                }

                @Override
                public void handleNack(long deliveryTag, boolean multiple) throws IOException {
                    if (multiple) {
                        // 如果开启了批量确认，则批量确认之后要清空有序集合已经确认好的元素
                        confirmSetNo.headSet(deliveryTag+1).clear();
                    } else {
                        // 如果没有开，那么只需要清除当前的这一个
                        confirmSetNo.remove(deliveryTag);
                    }
                    // 根据业务进行改进（重新发送消息）
                }
            });
            // 5.发送消息
            for (int i = 0; i < MESSAGE_COUNT; i++) {
                String msg = "hello publisher confirms " + i;
                // 拿到当前要发送的序号
                long seqNo = channel.getNextPublishSeqNo();
                channel.basicPublish("", Constants.PUBLISHER_CONFIRMS_QUEUE3, null, msg.getBytes());
                // 还未确认，放入集合
                confirmSetNo.add(seqNo);
            }
            while (!confirmSetNo.isEmpty()) {
                Thread.sleep(20);
            }
            long endTime = System.currentTimeMillis();
            System.out.printf("异步确认策略，消息条数：%d，耗时：%d ms \n", MESSAGE_COUNT, endTime - startTime);

        }
    }

    /**
     * 批量确认<br/>
     * 一次性确认一批
     */
    private static void publishingMessagesInBatches() throws Exception {
        try (Connection connection = createConnection()) {
            // 1.开启信道
            Channel channel = connection.createChannel();
            // 2.设置confirm模式
            channel.confirmSelect();
            // 3.使用内部交换机，声明队列
            channel.queueDeclare(Constants.PUBLISHER_CONFIRMS_QUEUE2, true, false, false, null);
            // 4.发送消息并返回确认
            int batchSize = 100;
            int outstandingMessages = 0;
            long startTime = System.currentTimeMillis();
            for (int i = 0; i < MESSAGE_COUNT; i++) {
                String msg = "hello publisher confirms " + i;
                channel.basicPublish("", Constants.PUBLISHER_CONFIRMS_QUEUE2, null, msg.getBytes());
                outstandingMessages++;
                // 批量确认
                if (outstandingMessages == batchSize) {
                    channel.waitForConfirms(5000);
                    outstandingMessages = 0;
                }
            }
            // 消息发送完毕之后如果有遗漏，进行确认
            if (outstandingMessages != 0) {
                channel.waitForConfirms(5000);
            }
            long endTime = System.currentTimeMillis();
            System.out.printf("批量确认策略，消息条数：%d，耗时：%d ms \n", MESSAGE_COUNT, endTime - startTime);
        }
    }

    /**
     * 单独确认<br/>
     * 每发送一个就要进行确认
     */
    private static void publishingMessagesIndividually() throws Exception {

        try (Connection connection = createConnection()) {
            // 1.开启信道
            Channel channel = connection.createChannel();
            // 2.设置confirm模式
            channel.confirmSelect();
            // 3.使用内部交换机，声明队列
            channel.queueDeclare(Constants.PUBLISHER_CONFIRMS_QUEUE1, true, false, false, null);
            // 4.发送消息并返回确认
            long startTime = System.currentTimeMillis();
            for (int i = 0; i < MESSAGE_COUNT; i++) {
                String msg = "hello publish confirms " + i;
                channel.basicPublish("", Constants.PUBLISHER_CONFIRMS_QUEUE1, null, msg.getBytes());
                // 等待确认
                channel.waitForConfirms(5000);
            }
            long endTime = System.currentTimeMillis();
            System.out.printf("单独确认策略，消息条数：%d，耗时：%d ms \n", MESSAGE_COUNT, endTime - startTime);
        }
    }
}
